package net.cyue.ort.llm.tokenizer;

import io.brunk.tokenizers.Encoding;
import net.cyue.util.BoxedUtil;
import net.cyue.util.FileUtil;
import net.cyue.util.ScalaUtil;
import scala.collection.Seq;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Path;

public class TransformersTokenizer extends Tokenizer {

    private io.brunk.tokenizers.Tokenizer tokenizer;

    public TransformersTokenizer(String configPath) {
        this.tokenizer = io.brunk.tokenizers.Tokenizer.fromFile(Path.of(configPath));
    }
    public TransformersTokenizer(File configFile) {
        this(configFile.getAbsolutePath());
    }
    public TransformersTokenizer(InputStream configStream) throws IOException {
        super(configStream);
        if (logger.isDebugEnabled()) {
            logger.debug("TransformersTokenizer initialized from InputStream");
        }
    }

    @Override
    protected void loadConfig(String configContent) throws IOException {
        String tmpDirPath = System.getProperty("java.io.tmpdir");
        File tmpFile = new File(new File(tmpDirPath), "tokenizer.json");
        FileUtil.writeStringToFile(tmpFile, configContent);
        this.tokenizer = io.brunk.tokenizers.Tokenizer.fromFile(
            Path.of(tmpFile.getAbsolutePath())
        );
        if (logger.isDebugEnabled()) {
            logger.debug(
                "Transformers tokenizer config loaded via temp file: {}, contentLength={}",
                tmpFile.getAbsolutePath(),
                configContent == null ? 0 : configContent.length()
            );
        }
    }

    @Override
    public String[] tokenize(String text) {
        long start = System.currentTimeMillis();
        Encoding encoding = this.tokenizer.encode(text, true);
        String[] tokens = ScalaUtil.seq2Array(encoding.tokens(), new String[0]);
        if (logger.isDebugEnabled()) {
            logger.debug(
                "tokenize() -> tokens={}, inputLength={}, time={}ms",
                tokens.length,
                text == null ? 0 : text.length(),
                System.currentTimeMillis() - start
            );
        }
        return tokens;
    }

    @Override
    public long[] tokens2IDs(String[] tokens) {
        long start = System.currentTimeMillis();
        Encoding encoding = this.tokenizer.encode(String.join("", tokens), true);
        long[] ids = ScalaUtil
                .seq2List(encoding.ids())
                .stream()
                .mapToLong(v-> Long.parseLong(v.toString()))
                .toArray();
        if (logger.isDebugEnabled()) {
            logger.debug(
                "tokens2IDs() -> ids={}, time={}ms",
            ids.length,
            System.currentTimeMillis() - start
            );
        }
        return ids;
    }

    @Override
    public Result encode(String text) {
        long start = System.currentTimeMillis();
        Encoding encoding = this.tokenizer.encode(text, true);
        String[] tokens = ScalaUtil.seq2Array(encoding.tokens(), new String[0]);
        Long[] ids = ScalaUtil.seq2Array(encoding.ids(), new Long[0]);
        Result result = new Result(text, tokens, BoxedUtil.unboxed(ids));
        if (logger.isDebugEnabled()) {
            logger.debug(
                "encode() -> tokens={}, ids={}, time={}ms",
                tokens.length,
                ids.length,
                System.currentTimeMillis() - start
            );
        }
        return result;
    }

    @Override
    public String decode(long[] tokenIds, boolean skipSpecialTokens) {
        long start = System.currentTimeMillis();
        Seq<Object> ids = ScalaUtil.array2Seq(BoxedUtil.boxed(tokenIds));
        String text = this.tokenizer.decode((scala.collection.immutable.Seq<Object>) ids, skipSpecialTokens);
        if (logger.isDebugEnabled()) {
            logger.debug(
                "decode() <- ids={}, skipSpecialTokens={}, outputLength={}, time={}ms",
                tokenIds == null ? 0 : tokenIds.length,
                skipSpecialTokens,
                text == null ? 0 : text.length(),
                System.currentTimeMillis() - start
            );
        }
        return text;
    }
}
